import sys
import cv2

"""

本程序实现了，从一张图片中提取黄色像素的区域，即提取出输油管道，为后续从输油管道上漏油斑点的检测做准备

"""

class ImageViewer:
    def __init__(self):
        super().__init__()

        self.initUI()

    def initUI(self):
        image_path = "oil_2.png"
        img = cv2.imread(image_path)
        img = cv2.resize(img, (640, 400))

        self.black_pixel_count = 0

        if img is not None:

            cv2.imshow("Step", img)
            cv2.waitKey(0)

            # 将图像从BGR转换到HSV
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

            # 定义黄色的HSV范围
            lower_yellow = (20, 100, 100)
            upper_yellow = (30, 255, 255)

            # 进行颜色阈值处理，得到黄色区域的掩码
            mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
            cv2.imshow("Step", mask_yellow)
            cv2.waitKey(0)

            # 查找黄色区域内的轮廓
            contours, _ = cv2.findContours(mask_yellow, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # 创建一个和原图像大小相同的全白图像
            result = img.copy()
            result.fill(255)

            # 遍历轮廓，将轮廓内的区域从原图像复制到结果图像
            for contour in contours:
                x, y, w, h = cv2.boundingRect(contour)
                #抠出黄色连通域
                roi = img[y:y + h, x:x + w]

                # 在黄色roi区域中，找到黑色区域
                lower_black = (0, 0, 0)
                upper_black = (180, 255, 30)
                mask_black = cv2.inRange(roi, lower_black, upper_black)
                cv2.imshow("Step", mask_black)
                cv2.waitKey(0)

                # 查找黑色区域内的轮廓
                black_contours, _ = cv2.findContours(mask_black, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) 

                # 在结果图像上，用红色画笔，画出黑色轮廓
                cv2.drawContours(roi, black_contours, -1, (0, 0, 255), 2)
                cv2.imshow("Step", roi)
                cv2.waitKey(0)

                # 遍历黑色轮廓，将轮廓内的区域从原图像复制到结果图像
                for black_contour in black_contours:
                    black_x, black_y, black_w, black_h = cv2.boundingRect(black_contour)
                    black_roi = roi[black_y:black_y + black_h, black_x:black_x + black_w]
                    result[y + black_y: y + black_y + black_h, x + black_x: x + black_x + black_w] = black_roi
                    # 计算轮廓内的黑色像素数量
                    contour_pixel_count = cv2.contourArea(contour)
                    self.black_pixel_count += contour_pixel_count
            #如果黄色区域内的黑色超过阈值，发生漏油，打出Leak字样
            if(self.black_pixel_count > 100):
            # 在结果图像上写出"Leak"字样
                cv2.putText(result, "Leak", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            # 显示结果图像
            cv2.imshow("Step", result)
            cv2.waitKey(0)

            cv2.imshow("Step", img)
            cv2.waitKey(0)            


if __name__ == '__main__':
    ex = ImageViewer()